from typing import Tuple
import numpy as np
import cv2


def white_bg(h: int, w: int) -> np.ndarray:
    """
    生成全白图像。
    Args:
        h (int): 图片的高度H
        w (int): 图片的高度W
    Returns:
        img (np.ndarray): 全白图像
    """
    return np.full((h, w, 3), 255, dtype = "uint8")


def black_bg(h: int, w: int) -> np.ndarray:
    """
    生成全黑图像。
    Args:
        h (int): 图片的高度H
        w (int): 图片的高度W
    Returns:
        img (np.ndarray): 全黑图像
    """
    return np.zeros((h, w, 3), dtype = "uint8")


def bgr2hex(bgr: np.ndarray) -> str:
    """
    将BGR颜色转成十六进制颜色字符串。
    Args:
        bgr (np.ndarray): BGR颜色向量
    Returns:
        hex (str): 十六进制颜色字符串
    """
    bgr = bgr.flatten()
    b, g, r = bgr[0], bgr[1], bgr[2]
    res = (r << 16) + (g << 8) + b
    res = hex(res).replace("0x", "")
    res = "0" * (6 - len(res)) + res
    return "#" + res


def hex2bgr(hex: str) -> np.ndarray:
    """
    将十六进制颜色字符串转成BGR颜色。
    Args:
        hex (str): 十六进制颜色字符串
    Returns:
        bgr (np.ndarray): BGR颜色向量
    """
    hex = int(hex.replace("0x", "").replace("#", ""), base = 16)
    bgr = np.array([[[hex & 0xff, (hex >> 8) & 0xff, (hex >> 16) & 0xff]]], dtype = "uint8")
    return bgr


def bgr2hsv(bgr: np.ndarray) -> np.ndarray:
    """
    将BGR颜色转成HSV颜色。
    Args:
        bgr (np.ndarray): BGR颜色向量
    Returns:
        hsv (np.ndarray): HSV颜色向量
    """
    hsv = cv2.cvtColor(bgr.astype("uint8"), cv2.COLOR_BGR2HSV)
    return hsv


def hsv2bgr(hsv: np.ndarray) -> np.ndarray:
    """
    将HSV颜色转成BGR颜色。
    Args:
        hsv (np.ndarray): HSV颜色向量
    Returns:
        bgr (np.ndarray): BGR颜色向量
    """
    bgr = cv2.cvtColor(hsv.astype("uint8"), cv2.COLOR_HSV2BGR)
    return bgr


def hsv2hex(hsv: np.ndarray) -> str:
    """
    将HSV颜色转成十六进制颜色字符串。
    Args:
        hsv (np.ndarray): HSV颜色向量
    Returns:
        hex (str): 十六进制颜色字符串
    """
    return bgr2hex(hsv2bgr(hsv))


def hex2hsv(hex: str) -> np.ndarray:
    """
    将十六进制颜色字符串转成HSV颜色。
    Args:
        hex (str): 十六进制颜色字符串
    Returns:
        hsv (np.ndarray): HSV颜色向量
    """
    return bgr2hsv(hex2bgr(hex))


def add_events(img: np.ndarray, vals: np.ndarray, y: np.ndarray, x: np.ndarray) -> np.ndarray:
    """
    添加事件。
    Args:
        img (np.ndarray): 空白图片，形状为[H, W, 3]
        vals (np.ndarray): 事件的值，形状为[N]，值域为[0, 1]，0为红，1为蓝，中间为红->蓝的过渡
        y (np.ndarray): 事件的纵坐标，形状为[N]，值域为[0, H - 1]
        x (np.ndarray): 事件的横坐标，形状为[N]，值域为[0, W - 1]
    Returns:
        rendered_img (np.ndarray): 渲染后的图片，形状为[H, W, 3]
    """
    get_color = lambda v: hsv2bgr(hex2hsv("#ff0000") * (1.0 - v) + hex2hsv("#0000ff") * v)
    y = img.shape[0] - 1 - y
    x = x.astype("int32")
    y = y.astype("int32")
    mask = (y >= 0) & (x < img.shape[1])
    vals = vals[mask]
    x = x[mask]
    y = y[mask]
    colors = [get_color(v) for v in vals]
    for i in range(len(colors)):
        cv2.circle(img, (x[i], y[i]), 1, tuple([int(v) for v in colors[i].flatten()]), -1)
    return img


def add_obstacle(img: np.ndarray, sizes: np.ndarray, y: np.ndarray, x: np.ndarray) -> np.ndarray:
    """
    添加障碍物。
    Args:
        img (np.ndarray): 空白图片，形状为[H, W, 3]
        sizes (np.ndarray): 障碍物的大小，形状为[N]
        y (np.ndarray): 障碍物的纵坐标，形状为[N]，值域为[0, H - 1]
        x (np.ndarray): 事件的纵坐标，形状为[N]，值域为[0, W - 1]
    Returns:
        rendered_img (np.ndarray): 渲染后的图片，形状为[H, W, 3]
    """
    get_color = lambda v: hsv2bgr(hex2hsv("#ff0000") * (1.0 - v) + hex2hsv("#0000ff") * v)
    y = img.shape[0] - 1 - y
    x = x.astype("int32")
    y = y.astype("int32")
    mask = (y >= 0) & (x < img.shape[1])
    x = x[mask]
    y = y[mask]
    sizes = sizes[mask]
    colors = [get_color(v) for v in np.random.random(sizes.shape[0])]
    for i in range(len(colors)):
        cv2.circle(img, (x[i], y[i]), 1, tuple([int(v) for v in colors[i].flatten()]), -1)
    return img


if __name__ == "__main__":
    x = np.random.randint(1280, size = (1000,))
    y = np.random.randint(800, size = (1000,))
    vals = np.random.randint(2, size = (1000,))
    img = white_bg(800, 1280)
    img = add_events(img, vals, y, x)
    cv2.imshow("img", img)
    cv2.waitKey(0)